import cv2
from configure import *
import shelve
import numpy as np
import os


class easySample():
    """"
        通过MD5值快速得到各个模型所需要的样本输入
    """
    sample = None

    def __init__(self) -> None:
        pass

    def get_sample(self, name, model):
        if model == "tlsh":
            return self.get_tlsh_sample(name)
        elif model == "siamese_img":
            return self.get_siamese_img_sample(name)
        elif model == "functionSim" or model == "RGCN":
            return self.get_functionSim_sample(name)
        elif model == "MGMN":
            return self.get_MGMN_sample(name)
        elif model == "siamese_graphsage":
            return self.get_siamese_graphsage_sample(name)
        elif model in ["angr", "retdec", "radare2", "ida"]:
            return self.get_disassemble_sample(name, model)
        elif model == "EECG":
            return self.get_EECG_sample(name)
        # 新增加了数据集
        elif model == "functionSim2" or model == "RGCN2":
            return self.get_functionSim2_sample(name)
        elif model == "siamese_img":
            return self.get_siamese_img2_sample(name)
        elif model == "MGMN":
            return self.get_MGMN2_sample(name)
        elif model == "siamese_graphsage2":
            return self.get_siamese_graphsage2_sample(name)

    def get_disassemble_sample(self, name, model):
        """
            获得不同工具处理后的样本
            格式一致，只需要换一下路径
        """
        dict = {}
        dict["angr"] = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/angr"
        dict["retdec"] = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/retdec"
        dict["radare2"] = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/radare2_functionSim"
        dict["ida"] = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/functionSim"
        return self.get_functionSim_sample(name, dict[model])

    def get_MGMN_sample(self, name):
        fileroot = functionSim_predata+"//"+name
        res = {}
        res["adj"], res["att"] = [], []
        with shelve.open(fileroot) as file:
            cg = file["cg"]
            cgattr = file["cgattr"]
            caller = file["caller"]  # 同一行是谁调用了我
            functype = file["funcType"]
            funcs_id = file["func_id"]  # name-->ind
        id_funcs = {}
        for i in funcs_id.keys():
            id_funcs[funcs_id[i]] = i

        #   动态导入函数是[],需要处理一下,
        #   同时生成vtype
        # tempData=[]
        for i in functype.keys():
            # temp=[]
            # if functype[i]=="local":
            # temp=[1,0,0]
            # elif functype[i]=='dynamic import':
            if functype[i] == 'dynamic import':
                # temp=[0,1,0]
                cgattr[i] = self.map_api_name_to_eight_bit_vector(
                    id_funcs[i], 8, 20)
            # else:
                # temp=[0,0,1]
            # tempData.append(temp)
        res["adj"] = np.array(cg)  # 这里是带边权的
        res["att"] = np.array(cgattr)  # 应该是8维，但是第一维字符串的提取似乎出现了问题，建议只是用后七维
        # res["vtype"]=np.array(tempData)
        return res

    def get_MGMN2_sample(self, name):
        fileroot = "/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/HGMSim"+"//"+name
        res = {}
        res["adj"], res["att"] = [], []
        with shelve.open(fileroot) as file:
            cg = file["cg"]
            cgattr = file["cgattr"]
            caller = file["caller"]  # 同一行是谁调用了我
            functype = file["funcType"]
            funcs_id = file["func_id"]  # name-->ind
        id_funcs = {}
        for i in funcs_id.keys():
            id_funcs[funcs_id[i]] = i

        for i in functype.keys():
            if functype[i] == 'dynamic import':
                cgattr[i] = self.map_api_name_to_eight_bit_vector(
                    id_funcs[i], 8, 20)
        res["adj"] = np.array(cg)  # 这里是带边权的
        res["att"] = np.array(cgattr)  # 应该是8维，但是第一维字符串的提取似乎出现了问题，建议只是用后七维
        return res

    def get_siamese_graphsage_sample(self, name):
        """
            ***将函数的汇编代码通过asm2vec模型转换成特征向量***
            返回图结构和图节点的汇编代码

        """
        radara2SamplePath = "/mnt/mydisk1/chenyongwei/sampleDatas/radare2_predata"

        fileroot = radara2SamplePath+name
        res = {}
        res["adj"], res["att"] = [], []
        # 增加判断文件是否存在，不然shelve会生成一个默认的
        if os.path.exists(fileroot+".dir"):
            with shelve.open(fileroot) as file:
                sampleInf = file["data"]
                # size:样本中函数的个数
                # discode:函数名及其对应的反汇编代码
                # adj:邻接矩阵
                # calls：函数间调用关系，都是函数名
                # name_to_id:函数名与下标的映射
                res["adj"] = sampleInf["adj"]
                res["name_to_id"] = sampleInf["name_to_id"]
                res["adj"] = sampleInf["adj"]
                res["att"] = sampleInf["att"]
        else:
            raise ValueError("样本embedding不存在!!!")
        return res

    def get_siamese_graphsage2_sample(self, name):
        """
            ***将函数的汇编代码通过asm2vec模型转换成特征向量***
            返回图结构和图节点的汇编代码

        """
        radara2SamplePath = "/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/radare2_predata/"

        fileroot = radara2SamplePath+name
        res = {}
        res["adj"], res["att"] = [], []
        if os.path.exists(fileroot+".dir"):
            with shelve.open(fileroot) as file:
                sampleInf = file["data"]
                res["adj"] = sampleInf["adj"]
                res["name_to_id"] = sampleInf["name_to_id"]
                res["adj"] = sampleInf["adj"]
                res["att"] = sampleInf["att"]
        else:
            raise ValueError("样本embedding不存在!!!")
        return res

    def get_functionSim_sample(self, name, filePath=r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/functionSim"):
        # def get_functionSim_sample(self,name,filePath=r"/mnt/mydisk1/chenyongwei/sampleDatas/functionSim"):
        # def get_functionSim_sample(self,name,filePath=r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/HGMSim"):
        # def get_functionSim_sample(self,name,filePath=r"/mnt/mydisk1/chenyongwei/sampleDatas/EECG"):
        """
            路径输入默认是ida进行反汇编
            输出：
                adj 邻接矩阵 b*n*n 
                att 特征矩阵 b*n*d
                vtype 类型矩阵 b*n*3  3种类型
        """
        fileroot = filePath+"//"+name
        res = {}
        res["adj"], res["att"], res["vtype"] = [], [], []
        with shelve.open(fileroot) as file:
            cg = file["cg"]
            cgattr = file["cgattr"]
            # caller=file["caller"]#同一行是谁调用了我
            # 这里的名称似乎写错了。。。
            try:
                functype = file["funcType"]
            except Exception as e:
                functype = file["functype"]
            try:
                funcs_id = file["func_id"]  # name-->ind
            except Exception as e:
                funcs_id = file["funcs_id"]  # name-->ind
        id_funcs = {}
        for i in funcs_id.keys():
            id_funcs[funcs_id[i]] = i

        #   动态导入函数是[],需要处理一下,
        #   同时生成vtype
        tempData = []
        for i in functype.keys():
            temp = []
            if functype[i] == "local":
                temp = [1, 0, 0]
            elif functype[i] == 'dynamic import':
                temp = [0, 1, 0]
                cgattr[i] = self.map_api_name_to_eight_bit_vector(
                    id_funcs[i], 8, 20)
            else:
                temp = [0, 0, 1]
            tempData.append(temp)
        res["adj"] = np.array(cg)  # 这里是带边权的
        res["att"] = np.array(cgattr)  # 应该是8维，但是第一维字符串的提取似乎出现了问题，建议只是用后七维
        res["vtype"] = np.array(tempData)
        return res

    def get_functionSim2_sample(self, name, filePath=r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/HGMSim"):
        """
            路径输入默认是ida进行反汇编
            输出：
                adj 邻接矩阵 b*n*n 
                att 特征矩阵 b*n*d
                vtype 类型矩阵 b*n*3  3种类型
        """
        fileroot = filePath+"//"+name
        res = {}
        res["adj"], res["att"], res["vtype"] = [], [], []
        with shelve.open(fileroot) as file:
            cg = file["cg"]
            cgattr = file["cgattr"]
            # caller=file["caller"]#同一行是谁调用了我
            # 这里的名称似乎写错了。。。
            try:
                functype = file["funcType"]
            except Exception as e:
                functype = file["functype"]
            try:
                funcs_id = file["func_id"]  # name-->ind
            except Exception as e:
                funcs_id = file["funcs_id"]  # name-->ind
        id_funcs = {}
        for i in funcs_id.keys():
            id_funcs[funcs_id[i]] = i

        #   动态导入函数是[],需要处理一下,
        #   同时生成vtype
        tempData = []
        for i in functype.keys():
            temp = []
            if functype[i] == "local":
                temp = [1, 0, 0]
            elif functype[i] == 'dynamic import':
                temp = [0, 1, 0]
                cgattr[i] = self.map_api_name_to_eight_bit_vector(
                    id_funcs[i], 8, 20)
            else:
                temp = [0, 0, 1]
            tempData.append(temp)
        res["adj"] = np.array(cg)  # 这里是带边权的
        res["att"] = np.array(cgattr)  # 应该是8维，但是第一维字符串的提取似乎出现了问题，建议只是用后七维
        res["vtype"] = np.array(tempData)
        return res

    def get_EECG_sample(self, name, filePath=r"/mnt/mydisk1/chenyongwei/sampleDatas/EECG"):
        """
            路径输入默认是ida进行反汇编
            输出：
                adj 邻接矩阵 b*n*n 
                att 特征矩阵 b*n*d
                vtype 类型矩阵 b*n*3  3种类型
        """
        fileroot = filePath+"//"+name
        res = {}
        res["adj"], res["att"], res["vtype"] = [], [], []

        with shelve.open(fileroot) as file:
            res["adj"] = file['adj']
            res["att"] = file['att']
            res["vtype"] = file['vType']
        return res

    def get_tlsh_sample(self, name):
        return name

    def get_siamese_img_sample(self, name):
        # imgPath = "/home/cyw/projects/function_sim_project/all_data/sampleDatas/img"
        imgPath = "/mnt/mydisk1/chenyongwei/sampleDatas/img"
        res = cv2.imread(imgPath+"/"+name+".png", cv2.IMREAD_UNCHANGED)
        return res

    def get_siamese_img2_sample(self, name):
        imgPath = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/img"
        res = cv2.imread(imgPath+"/"+name+".png", cv2.IMREAD_UNCHANGED)
        return res

    def map_api_name_to_eight_bit_vector(self, function_name, embeddingSize, numLim):
        """
            将api转换成一个embeddingSize位的向量，每一位使用numLim取余
        """
        function_bytes = function_name.encode('utf-8')
        eight_bit_vector = [0] * embeddingSize
        for i, byte in enumerate(function_bytes):
            index = i % embeddingSize
            eight_bit_vector[index] += byte
            eight_bit_vector[index] %= numLim
        return eight_bit_vector


if __name__ == "__main__":
    a = easySample()
    b = a.get_sample("0a0c8e3a001ad2c11c95602c6fe4b537", "functionSim2")
    print(b)
